﻿#ifndef DNNDETECTOR_H
#define DNNDETECTOR_H
#include <vector>
#include <string>
#include <opencv2/imgproc.hpp>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
#include <random>

struct Detection
{
    int class_id{0};
    std::string className{};
    float confidence{0.0};
    cv::Scalar color{};
    cv::Rect box{};
};

class DnnDetector
{
public:
    DnnDetector(const std::string &onnxModelPath, const cv::Size &modelInputShape = {640, 640}, std::vector<std::string> classNames={}, const bool &runWithCuda = true);

    std::vector<Detection> detect(const cv::Mat &input);


private:
    std::vector<std::string> classes;
    cv::dnn::Net net;
    std::string modelPath{};
    bool cudaEnabled{};
    cv::Size2f modelShape{};
    bool letterBoxForSquare = true;
    float modelConfidenceThreshold {0.25};
    float modelScoreThreshold      {0.45};
    float modelNMSThreshold        {0.50};


    void loadOnnxNetwork();
    cv::Mat formatToSquare(const cv::Mat &source);


};

#endif // DNNDETECTOR_H
